library(metafor)
library(plyr)
library(tidyverse)
library(magrittr)
library(foreign)
library(estimatr)
library(emmeans)
library(broom)
library(countrycode)

t4 <- "https://github.com/thomasjwood/kz_multicountry/raw/main/t3.rds" %>% 
  url %>% 
  gzcon %>% 
  readRDS

t5 <- t4 %>% 
  select(
    -(vaxskep_safe:vaxskep_child), 
    -(crt_ball:crt_lily),
    -(consp_cabal:consp_cures),
    -ends_with("_scale")
  ) %>% 
  modify_at(
    str_c(
      c("vax", "consp", "crt"), "_grp"
    ),
    ~factor(
      ., 
      c("Low", "Med", "High")
    )
  ) %>% 
  gather(
    scale, val, ends_with("_grp")
  ) %>% 
  arrange(
    geo_country_code, caseid
  ) %>% 
  group_by(
    geo_country_code, wave, issue, scale, val
  ) %>% 
  nest %>% 
  ungroup %>% 
  mutate(
    val = val %>% 
      factor(
        c("Low", "Med", "High")
      )
  ) %>% 
  arrange(geo_country_code, issue, scale, val) %>% 
  filter(
    wave == "wave 1"
  ) 

t5 <- t4 %>% 
  select(
    -(vaxskep_safe:vaxskep_child), 
    -(crt_ball:crt_lily),
    -(consp_cabal:consp_cures),
    -ends_with("_scale")
  ) %>% 
  modify_at(
    str_c(
      c("vax", "consp", "crt"), "_grp"
    ),
    ~factor(
      ., 
      c("Low", "Med", "High")
    )
  ) %>% 
  gather(
    scale, val, ends_with("_grp")
  ) %>% 
  arrange(
    geo_country_code, caseid
  ) %>% 
  group_by(
    geo_country_code, wave, issue, scale, val
  ) %>% 
  nest %>% 
  ungroup %>% 
  mutate(
    val = val %>% 
      factor(
        c("Low", "Med", "High")
      )
  ) %>% 
  arrange(geo_country_code, issue, scale, val) %>% 
  filter(
    wave == "wave 1"
  )


t5$mods <- t5$data %>% 
  map(
    \(i)
    
    lm_robust(
      out_num ~ treatment,
      data = i
    )
  )

t5$emm <- t5$mods %>% 
  map(
    ~emmeans(
      ., 
      consec ~ treatment, 
      adjust = "none"
    )
  )


t5_2 <- t5 %>%
  ungroup %>% 
  select(geo_country_code, scale, val, emm) %>% 
  pmap_dfr(
    function(geo_country_code, scale, val, emm)
      emm$contrasts %>% 
      tidy %>% 
      mutate(
        geo_country_code, scale, val,
        wave = "wave 1",
        geo_country_code = "meta"
      )
  ) %>% 
  nest_by(
    contrast, scale, val
  )

t5_2$rma <- t5_2$data %>% 
  map(
    \(i)
    rma(yi = i$estimate, sei = i$std.error)
  )

t6 <- t5 %>% 
  ungroup %>% 
  select(
    geo_country_code, wave, issue, scale, val, emm
  ) %>% 
  pmap_dfr(
    \(geo_country_code, wave, issue, scale, val, emm)
    emm %>% 
      map_dfr(
        ~tidy(.) %>%
          rename_with(
            ~str_replace(., "contrast", "treatment")
          ),
        .id = "type"  
      ) %>% 
      select(-term, -null.value) %>% 
      mutate(
        geo_country_code, wave, issue, scale, val
      )
  ) %>% 
  mutate(
    cvar = case_when(
      type %>% 
        equals("emmeans") ~ "Conditional means",
      type %>% 
        equals("emmeans") %>% 
        not  &
        treatment %>% 
        str_detect("Misinformation - ") ~ "Misinformation effect",
      type %>% 
        equals("emmeans") %>% 
        not  &
        treatment %>% 
        str_detect("Misinformation &") ~ "Correction effect"
    ) %>% 
      factor(
        c("Conditional means",
          "Misinformation effect",
          "Correction effect")
      )
  ) %>% 
  filter(
    type == "contrasts"
  ) %>% 
  bind_rows(
    t5_2 %>%
      ungroup %>% 
      pmap_dfr(
        function(contrast,  scale, val, data, rma)
          
          rma %>%
          tidy %>%
          mutate(
            contrast, scale, val, 
            wave = data$wave[[1]],
            geo_country_code = data$geo_country_code[[1]]
          )
      ) %>%
      mutate(
        cvar = case_when(
          contrast %>%
            str_detect("Misinformation &") ~ "Correction effect",
          TRUE ~ "Misinformation effect"
        )  %>%
          factor(
            c("Conditional means",
              "Misinformation effect",
              "Correction effect")
          ),
        type = "meta"
      ) %>%
      rename(
        treatment = contrast
      ) %>%
      select(-term)
  ) 


t6$labs <- t6$estimate %>%
  round(2) %>% 
  abs %>% 
  equals(0) %>% 
  ifelse(
    t6$estimate %>% 
      round(3) %>%
      str_replace("0\\.", "."),
    t6$estimate %>% 
      round(2) %>%
      str_replace("0\\.", ".")
  )

t6$labs %<>% 
  str_detect("-") %>%
  ifelse(
    t6$labs %>% 
      str_pad(width = 4, side = "right", pad = "0"),
    t6$labs %>% 
      str_pad(width = 3, side = "right", pad = "0")
  )

t6$labs[t6$labs == "000"] <- t6$estimate[t6$labs == "000"] %>% 
  round(4) %>%
  format(scientific = F) %>% 
  str_replace("0\\.", ".")

t6$labs %<>% 
  str_c(
    t6$p.value %>% 
      gtools::stars.pval() %>% 
      str_remove("\\.")
  )


t6$low <- t6$estimate %>% 
  subtract(
    t6$std.error %>% 
      multiply_by(1.96)
  )

t6$hi <- t6$estimate %>% 
  add(
    t6$std.error %>% 
      multiply_by(1.96)
  )

t6$issue %<>%
  mapvalues(
    c("global 1",
      "global 2",
      "country specific",
      NA),
    c("The Gates Foundation is going to make huge profits from COVID-19 vaccines.",
      "mRNA vaccines like Moderna permanently alter DNA, according to Moderna executives.",
      "Country specific misinformation",
      "Overall") %>%
      str_wrap(17)
  )

t6$stat_grp <- t6$p.value %>% 
  is_less_than(.05) %>% 
  ifelse(
    "p < .05", "p >= .05"
  )

t6$treatment %<>%
  mapvalues(
    t6$treatment %>% 
      unique,
    c("Misinformation effect",
      "Correction effect")
  ) %>% 
  factor(
    c("Misinformation effect",
      "Correction effect")
  )

t6$scale %<>% 
  mapvalues(
    t6$scale %>% 
      unique,
    c("Conspiracy scale",
      "Cognitive Reflection Test",
      "Vaccine skepticism scale")
  )

library(ggbeeswarm)

p1 <- t6 %>%
  filter(
    type  %>% 
      equals("meta") %>% 
      not
  ) %>% 
  ggplot() +
  geom_hline(
    yintercept =  0,
    linetype = "dashed",
    size = .375
  ) +
  geom_label(
    aes(
      val %>% factor, estimate, label = labs
    ),
    data = t6 %>% 
      filter(
        type %>% 
          equals("meta") 
      ),
    fill = "grey95",
    label.size =  0,
    nudge_x = .35, 
    size = 2.5,
  ) +
  geom_dotplot(
    aes(
      x = val %>% factor, y = estimate, fill = stat_grp
    ),
    stackdir = "center",
    stackratio = 1,
    dotsize = 1.25,
    
    binaxis = 'y',
    binwidth = 1.2/50,
    data = t6 %>% 
      filter(
        type %>% 
          equals("meta") %>% 
          not
      )
  ) +
  coord_flip() +
  facet_grid(
    scale ~ treatment, 
    scales = "free", 
    space = "free",
    labeller = label_wrap_gen(width = 15)
  ) +
  geom_linerange(
    aes(
      ymin = low, ymax = hi, val %>% factor
    ),
    data = t6 %>% 
      filter(
        type %>% 
          equals("meta") 
      ),
    color = "white",
    size = 1.4
  ) +
  geom_linerange(
    aes(
      ymin = low, ymax = hi, val %>% factor
    ),
    data = t6 %>% 
      filter(
        type %>% 
          equals("meta") 
      ),
    size = .3
  ) +
  geom_point(
    aes(
      val %>% factor, estimate
    ),
    data = t6 %>% 
      filter(
        type %>% 
          equals("meta") 
      ),
    size = 2,
    shape = 23,
    fill = "white"
  ) +
  geom_segment(
    aes(
      x = x, y = y, xend = xend, yend = yend
    ),
    arrow = arrow(length = unit(0.2, "cm"), type = "closed"),
    data = tribble(
      ~x, ~y, ~xend, ~yend,
      .5, -.02, .5, -.4,
      .5,  .02, .5,  .4
    ) %>% 
      mutate(
        treatment = t6$treatment[[1]] %>% 
          factor(
            t6$treatment %>% 
              levels
          ),
        scale = t6$scale[[13]] %>% 
          factor(
            t6$scale %>% 
              factor %>% 
              levels
          )
      ),
    size = .25
  ) +
  geom_text(
    aes(
      x, y, label = lab
    ),
    data = tribble(
      ~x,  ~y, ~lab,
      .6,   .19, "Less accurate",
      .6, -.19, "More accurate",
      
    ) %>% 
      mutate(
        treatment = t6$treatment[[1]] %>% 
          factor(
            t6$treatment %>% 
              levels
          ),
        scale = t6$scale[[13]] %>% 
          factor(
            t6$scale %>% 
              factor %>% 
              levels
          )
      ), 
    size = 2,
    family = "Roboto", 
    fontface = "italic"
  ) +
  labs(
    x = "Pretreatment psychological groups",
    fill = "Statistical significance",
    y = "Difference on four point scale"
  ) +
  scale_fill_grey(
    start = .4, end = .99,
    guide = guide_legend(
      direction = "horizontal", 
      title.position = "left"
    )
  ) +
  scale_y_continuous(
    breaks = seq(-.6, .9, .3),
    labels = c(
      "-.6", "-.3", "0", ".3", ".6", ".9"
    )
  )
